from ultralytics import YOLO

class YOLOService:
    def __init__(self):
        self.model = None
        
    def get_model(self):
        """获取或初始化YOLO模型"""
        if self.model is None:
            self.model = YOLO('yolov8s.pt')
        return self.model
        
    def detect(self, image_path, class_mapping):
        """执行目标检测"""
        model = self.get_model()
        results = model(image_path)
        
        # 获取类别名称（英文）
        english_class_names = [model.names[int(c)] for c in results[0].boxes.cls]
        # 转换为中文类别名称
        chinese_class_names = [class_mapping.get(name, name) for name in english_class_names]
        
        # 构建检测数据
        detection_info = {
            'boxes': results[0].boxes.xyxy.tolist(),
            'confidences': results[0].boxes.conf.tolist(),
            'class_ids': results[0].boxes.cls.tolist(),
            'class_names': chinese_class_names,
            'original_names': english_class_names
        }
        
        return results[0], detection_info 